import numpy as np
import torch
import pickle as pkl
import random
import os
from deepctr_torch.inputs import SparseFeat, DenseFeat, VarLenSparseFeat, get_feature_names
from deepctr_torch.models import DIN

MAX_LEN = 2

def seed_torch(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
seed_torch(1234)


keys = ['customer_id', 'FN', 'Active', 'club_member_status',\
     'fashion_news_frequency', 'age', 'artid','pcode', 'y']

feature_columns = [SparseFeat('customer_id', 1371980, embedding_dim=64, use_hash=False),
                   SparseFeat('FN', 2, embedding_dim=4, use_hash=False),
                   SparseFeat('Active', 2, embedding_dim=4, use_hash=False),
                   SparseFeat('club_member_status', 2, embedding_dim=4, use_hash=False),
                   SparseFeat('fashion_news_frequency', 2, embedding_dim=4, use_hash=False),
                   SparseFeat('artid', 105542 + 1, embedding_dim=64, use_hash=False),
                   SparseFeat('pcode', 47224 + 1, embedding_dim=64, use_hash=False),
                   DenseFeat('age', 1)]
feature_columns +=[VarLenSparseFeat(SparseFeat('hist_artid', vocabulary_size=105542 + 1, embedding_dim=64, embedding_name='artid'), maxlen=MAX_LEN, length_name="seq_length"),
                   VarLenSparseFeat(SparseFeat('hist_pcode', vocabulary_size=47224 + 1, embedding_dim=64, embedding_name='pcode'), maxlen=MAX_LEN,length_name="seq_length")]
behavior_feature_list = ["artid", "pcode"]

with open('./test.pkl', 'rb') as f:
    test_set = pkl.load(f)
len_test = len(test_set)

test_set = list(zip(*test_set))

test_dict = {}
for i in range(len(keys)):
    test_dict[keys[i]] = np.asarray(test_set[i])
test_dict['seq_length'] = np.ones((len_test, )) * MAX_LEN
test_dict['hist_artid'] = np.ones((len_test, MAX_LEN))
test_dict['hist_pcode'] = np.ones((len_test, MAX_LEN))
X_test = {name: test_dict[name] for name in get_feature_names(feature_columns)}

device = 'cpu'
use_cuda = True
if use_cuda and torch.cuda.is_available():
    print('cuda ready...')
    device = 'cuda:4'

model = DIN(feature_columns, behavior_feature_list, device=device, att_weight_normalization=True, seed=1234)
model.load_state_dict(torch.load("./model0.pt"))
y_test = model.predict(X_test, 4096)
with open('./y_test.pkl', 'wb') as f:
    pkl.dump(y_test, f, pkl.HIGHEST_PROTOCOL)

